from PIL import Image, ImageFilter, ImageDraw, ImageFont
import os
import math
import random
import logging

os.chdir(os.path.dirname(os.path.abspath(__file__)))

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("gen_poster")

# 配置参数
TEMP_FOLDER = "temp"  # 海报来源文件夹
FONT_FOLDER = "fonts"  # 字体文件夹
CHINESE_FONT = "ch.ttf"  # 中文字体
ENGLISH_FONT = "en.otf"  # 英文字体

# 海报生成配置 - 基础1080p版本
POSTER_GEN_CONFIG = {
    "ROWS": 3,  # 行数
    "COLS": 3,  # 列数
    "MARGIN": 20,  # 图片间距
    "CORNER_RADIUS": 20,  # 圆角半径
    "ROTATION_ANGLE": -15,  # 旋转角度
    "START_X": 1000,  # 起始X坐标
    "START_Y": -240,  # 起始Y坐标
    "COLUMN_SPACING": 85,  # 列间距
    "SAVE_COLUMNS": False,  # 是否保存中间列图片
    "CELL_WIDTH": 350,  # 图片宽度
    "CELL_HEIGHT": 500,  # 图片高度
    "BACKGROUND_CORNER_RADIUS": 60,  # 整体背景圆角半径
}

def add_shadow(img, offset=(5, 5), shadow_color=(0, 0, 0, 100), blur_radius=3):
    """给图片添加右侧和底部阴影"""
    shadow_width = img.width + offset[0] + blur_radius * 2
    shadow_height = img.height + offset[1] + blur_radius * 2

    shadow = Image.new("RGBA", (shadow_width, shadow_height), (0, 0, 0, 0))
    shadow_layer = Image.new("RGBA", img.size, shadow_color)
    shadow.paste(shadow_layer, (blur_radius + offset[0], blur_radius + offset[1]))
    shadow = shadow.filter(ImageFilter.GaussianBlur(blur_radius))

    result = Image.new("RGBA", shadow.size, (0, 0, 0, 0))
    result.paste(img, (blur_radius, blur_radius), img if img.mode == "RGBA" else None)
    shadow_img = Image.alpha_composite(shadow, result)

    return shadow_img

def draw_text_on_image(
    image, text, position, font_path, font_size, fill_color=(255, 255, 255, 255), 
    shadow_enabled=False, shadow_color=(0, 0, 0, 180), shadow_offset=(2, 2)
):
    """在图像上绘制文字，可选添加文字阴影"""
    img_copy = image.copy()
    draw = ImageDraw.Draw(img_copy)
    
    try:
        font = ImageFont.truetype(font_path, font_size)
    except IOError:
        logger.warning(f"字体文件不存在:{font_path}，使用默认字体")
        font = ImageFont.load_default()
    
    # 如果启用阴影，先绘制阴影文字
    if shadow_enabled:
        shadow_position = (position[0] + shadow_offset[0], position[1] + shadow_offset[1])
        draw.text(shadow_position, text, font=font, fill=shadow_color)
    
    # 绘制正常文字
    draw.text(position, text, font=font, fill=fill_color)

    return img_copy

def draw_multiline_text_on_image(
    image, text, position, font_path, font_size, line_spacing=10,
    fill_color=(255, 255, 255, 255), shadow_enabled=False, 
    shadow_color=(0, 0, 0, 180), shadow_offset=(2, 2)
):
    """在图像上绘制多行文字，根据空格自动换行，可选添加文字阴影"""
    img_copy = image.copy()
    draw = ImageDraw.Draw(img_copy)
    
    try:
        font = ImageFont.truetype(font_path, font_size)
    except IOError:
        logger.warning(f"字体文件不存在:{font_path}，使用默认字体")
        font = ImageFont.load_default()

    # 按空格分割文本
    lines = text.split(" ")

    # 如果只有一行，直接绘制并返回
    if len(lines) <= 1:
        if shadow_enabled:
            shadow_position = (position[0] + shadow_offset[0], position[1] + shadow_offset[1])
            draw.text(shadow_position, text, font=font, fill=shadow_color)
        draw.text(position, text, font=font, fill=fill_color)
        return img_copy, 1

    # 绘制多行文本
    x, y = position
    for i, line in enumerate(lines):
        current_y = y + i * (font_size + line_spacing)
        
        # 如果启用阴影，先绘制阴影文字
        if shadow_enabled:
            shadow_x = x + shadow_offset[0]
            shadow_y = current_y + shadow_offset[1]
            draw.text((shadow_x, shadow_y), line, font=font, fill=shadow_color)
        
        # 绘制正常文字
        draw.text((x, current_y), line, font=font, fill=fill_color)

    # 返回图像和行数
    return img_copy, len(lines)

def get_random_color(image_path):
    """获取图片随机位置的颜色"""
    try:
        img = Image.open(image_path)
        width, height = img.size
        random_x = random.randint(int(width * 0.5), int(width * 0.8))
        random_y = random.randint(int(height * 0.5), int(height * 0.8))

        if img.mode == "RGBA":
            r, g, b, a = img.getpixel((random_x, random_y))
            return (r, g, b, a)
        elif img.mode == "RGB":
            r, g, b = img.getpixel((random_x, random_y))
            return (r + 100, g + 50, b, 255)
        else:
            img = img.convert("RGBA")
            r, g, b, a = img.getpixel((random_x, random_y))
            return (r, g, b, a)
    except Exception as e:
        logger.error(f"获取图片颜色时出错: {e}")
        return (random.randint(50, 200), random.randint(50, 200), random.randint(50, 200), 255)

def draw_color_block(image, position, size, color):
    """在图像上绘制色块"""
    img_copy = image.copy()
    draw = ImageDraw.Draw(img_copy)
    draw.rectangle([position, (position[0] + size[0], position[1] + size[1])], fill=color)
    return img_copy

def create_gradient_background(width, height, color=None):
    """创建一个从左到右的渐变背景"""
    def normalize_rgb(input_rgb):
        if isinstance(input_rgb, tuple):
            if len(input_rgb) == 2 and isinstance(input_rgb[0], tuple):
                return normalize_rgb(input_rgb[0])
            if len(input_rgb) == 4 and all(isinstance(v, (int, float)) for v in input_rgb):
                return input_rgb[:3]
            if len(input_rgb) == 3 and all(isinstance(v, (int, float)) for v in input_rgb):
                return input_rgb
        raise ValueError(f"无法识别的颜色格式: {input_rgb!r}")

    def is_mid_bright_hsl(input_rgb, min_l=0.3, max_l=0.7):
        r, g, b = normalize_rgb(input_rgb)
        r1, g1, b1 = r/255.0, g/255.0, b/255.0
        h, l, s = colorsys.rgb_to_hls(r1, g1, b1)
        return min_l <= l <= max_l
    
    import colorsys
    
    selected_color = None
    
    # 如果传入的是颜色数组
    if isinstance(color, list) and len(color) > 0:
        for i in range(min(10, len(color))):
            if is_mid_bright_hsl(color[i]):
                if isinstance(color[i], tuple) and len(color[i]) == 2 and isinstance(color[i][0], tuple):
                    selected_color = color[i][0]
                else:
                    selected_color = color[i]
                break
    
    # 如果没有找到合适的颜色，随机生成一个颜色
    if selected_color is None:
        def random_hsl_to_rgb(
            hue_range=(0, 360),
            sat_range=(0.5, 1.0),
            light_range=(0.5, 0.8)
        ):
            h = random.uniform(hue_range[0]/360.0, hue_range[1]/360.0)
            s = random.uniform(sat_range[0], sat_range[1])
            l = random.uniform(light_range[0], light_range[1])
            r, g, b = colorsys.hls_to_rgb(h, l, s)
            return (int(r*255), int(g*255), int(b*255))

        selected_color = random_hsl_to_rgb()

    # 左侧颜色 - 变亮
    r_left = min(255, int(selected_color[0] * 1.8))
    g_left = min(255, int(selected_color[1] * 1.8))
    b_left = min(255, int(selected_color[2] * 1.8))
    left_color = (r_left, g_left, b_left, selected_color[3] if len(selected_color) > 3 else 255)

    # 确保left_color包含alpha通道
    if len(left_color) == 3:
        left_color = (left_color[0], left_color[1], left_color[2], 255)
    
    # 右侧颜色 - 变暗
    r_right = int(selected_color[0] * 0.4)
    g_right = int(selected_color[1] * 0.4)
    b_right = int(selected_color[2] * 0.4)
    r_right = max(0, r_right)
    g_right = max(0, g_right)
    b_right = max(0, b_right)
    right_color = (r_right, g_right, b_right, left_color[3])
    
    # 创建左右两个纯色图像
    left_image = Image.new("RGBA", (width, height), left_color)
    right_image = Image.new("RGBA", (width, height), right_color)
    
    # 创建渐变遮罩 - 调整渐变曲线使左侧更亮
    mask = Image.new("L", (width, height), 0)
    mask_data = []
    
    for y in range(height):
        for x in range(width):
            # 使用幂函数调整渐变曲线，使左侧更亮，右侧更暗
            mask_value = int(255.0 * (x / width) ** 0.6)
            mask_data.append(mask_value)
    
    mask.putdata(mask_data)
    
    # 使用遮罩合成左右两个图像
    gradient = Image.composite(right_image, left_image, mask)
    
    return gradient

def get_poster_primary_color(image_path):
    """分析图片并提取主色调"""
    try:
        from collections import Counter
        
        img = Image.open(image_path)
        img = img.resize((100, 150), Image.LANCZOS)
        
        if img.mode != 'RGBA':
            img = img.convert('RGBA')
            
        pixels = list(img.getdata())
        
        # 过滤掉接近黑色和白色的像素，以及透明度低的像素
        filtered_pixels = []
        for pixel in pixels:
            r, g, b, a = pixel
            
            if a < 200:
                continue
                
            brightness = (r + g + b) / 3
            
            if brightness < 30 or brightness > 220:
                continue
                
            filtered_pixels.append((r, g, b, 255))
            
        if not filtered_pixels:
            filtered_pixels = [(p[0], p[1], p[2], 255) for p in pixels if p[3] > 100]
            
        if not filtered_pixels:
            return [(150, 100, 50, 255)]
            
        color_counter = Counter(filtered_pixels)
        common_colors = color_counter.most_common(10)
        
        if common_colors:
            return common_colors
        
        r_avg = sum(p[0] for p in filtered_pixels) // len(filtered_pixels)
        g_avg = sum(p[1] for p in filtered_pixels) // len(filtered_pixels)
        b_avg = sum(p[2] for p in filtered_pixels) // len(filtered_pixels)
        
        return [(r_avg, g_avg, b_avg, 255)]
     
    except Exception as e:
        logger.error(f"获取图片主色调时出错: {e}")
        return [(150, 100, 50, 255)]

def generate_poster(
    title="封面标题", 
    subtitle="封面副标题",
    output_filename="poster_output.png",
    resolution="720p"  # 新增参数：目标分辨率，支持'720p'、'480p'、'1080p'
):
    """
    生成海报封面，从temp文件夹获取9张海报图片
    
    参数:
        title: 封面主标题
        subtitle: 封面副标题
        output_filename: 输出文件名
        resolution: 输出分辨率，可选值：'1080p'、'720p'、'480p'
    """
    try:
        logger.info(f"正在生成{resolution}海报...")
        logger.info("-" * 40)

        # 根据分辨率设置目标尺寸
        if resolution == '1080p':
            target_width, target_height = 1920, 1080
            scale_factor = 1.0
        elif resolution == '720p':
            target_width, target_height = 1280, 720
            scale_factor = 1280 / 1920  # 约为0.6667
        elif resolution == '480p':
            target_width, target_height = 854, 480  # 使用480p的常见宽高比
            scale_factor = 854 / 1920  # 约为0.4448
        else:
            logger.error(f"不支持的分辨率: {resolution}，默认使用720p")
            target_width, target_height = 1280, 720
            scale_factor = 1280 / 1920
        
        # 直接使用输出文件夹和输出文件名拼接路径
        output_path = output_filename
        
        # 从配置中获取参数
        rows = POSTER_GEN_CONFIG["ROWS"]
        cols = POSTER_GEN_CONFIG["COLS"]
        margin = POSTER_GEN_CONFIG["MARGIN"]
        corner_radius = POSTER_GEN_CONFIG["CORNER_RADIUS"]
        rotation_angle = POSTER_GEN_CONFIG["ROTATION_ANGLE"]
        start_x = POSTER_GEN_CONFIG["START_X"]
        start_y = POSTER_GEN_CONFIG["START_Y"]
        column_spacing = POSTER_GEN_CONFIG["COLUMN_SPACING"]
        save_columns = POSTER_GEN_CONFIG["SAVE_COLUMNS"]
        cell_width = POSTER_GEN_CONFIG["CELL_WIDTH"]
        cell_height = POSTER_GEN_CONFIG["CELL_HEIGHT"]

        # 定义模板尺寸 - 始终使用1080p进行布局计算
        template_width = 1920
        template_height = 1080

        # 获取temp文件夹中的所有图片
        temp_folder = os.path.join(os.getcwd(), TEMP_FOLDER)
        if not os.path.exists(temp_folder):
            logger.error(f"错误: 找不到temp文件夹: {temp_folder}")
            return False

        # 支持的图片格式
        supported_formats = (".jpg", ".jpeg", ".png", ".bmp", ".gif", ".webp")
        
        # 获取并排序图片
        poster_files = [
            os.path.join(temp_folder, f)
            for f in os.listdir(temp_folder)
            if os.path.isfile(os.path.join(temp_folder, f))
            and f.lower().endswith(supported_formats)
        ]

        # 确保至少有一张图片
        if not poster_files:
            logger.error(f"错误: 在 {temp_folder} 中没有找到支持的图片文件")
            return False

        # 限制最多处理 rows*cols 张图片
        max_posters = rows * cols
        poster_files = poster_files[:max_posters]

        # 如果图片数量不足，使用循环方式填充
        if len(poster_files) < max_posters:
            # 计算还需要添加的图片数量
            needed = max_posters - len(poster_files)
    
            # 从原始图片列表中循环选择图片添加
            for i in range(needed):
                # 使用取模运算循环选择原始图片列表中的图片
                poster_files.append(poster_files[i % len(poster_files)])
    
            # 验证最终列表长度是否正确
            assert len(poster_files) == max_posters, f"海报数量应为 {max_posters}，但实际为 {len(poster_files)}"

        # 自定义排序顺序
        custom_order = "315426987"
        order_map = {num: index for index, num in enumerate(custom_order)}

        # 按自定义顺序排序图片
        poster_files = sorted(
            poster_files,
            key=lambda x: order_map.get(str(poster_files.index(x) + 1), 0)
        )

        # 获取第一张图片的主色调，用于生成背景
        color = get_poster_primary_color(poster_files[0])
        
        # 创建渐变背景
        gradient_bg = create_gradient_background(template_width, template_height, color)

        # 创建保存中间文件的文件夹
        columns_dir = os.path.join("columns")
        if save_columns and not os.path.exists(columns_dir):
            os.makedirs(columns_dir)

        # 将图片分成3组，每组3张
        grouped_posters = [
            poster_files[i : i + rows] for i in range(0, len(poster_files), rows)
        ]

        # 以渐变背景作为起点
        result = gradient_bg.copy()
        
        # 处理每一组（每一列）图片
        for col_index, column_posters in enumerate(grouped_posters):
            if col_index >= cols:
                break

            # 计算当前列的 x 坐标
            column_x = start_x + col_index * column_spacing

            # 计算当前列所有图片组合后的高度（包括间距）
            column_height = rows * cell_height + (rows - 1) * margin

            # 创建一个透明的画布用于当前列的所有图片
            shadow_extra_width = 20 + 20 * 2
            shadow_extra_height = 20 + 20 * 2

            column_image = Image.new(
                "RGBA",
                (cell_width + shadow_extra_width, column_height + shadow_extra_height),
                (0, 0, 0, 0),
            )

            # 在列画布上放置每张图片
            for row_index, poster_path in enumerate(column_posters):
                try:
                    # 打开海报
                    poster = Image.open(poster_path)

                    # 调整海报大小为固定尺寸
                    resized_poster = poster.resize(
                        (cell_width, cell_height), Image.LANCZOS
                    )

                    # 创建圆角遮罩
                    if corner_radius > 0:
                        mask = Image.new("L", (cell_width, cell_height), 0)
                        draw = ImageDraw.Draw(mask)
                        draw.rounded_rectangle(
                            [(0, 0), (cell_width, cell_height)],
                            radius=corner_radius,
                            fill=255,
                        )

                        poster_with_corners = Image.new(
                            "RGBA", resized_poster.size, (0, 0, 0, 0)
                        )
                        poster_with_corners.paste(resized_poster, (0, 0), mask)
                        resized_poster = poster_with_corners

                    # 添加阴影效果到每张海报
                    resized_poster_with_shadow = add_shadow(
                        resized_poster,
                        offset=(20, 20),
                        shadow_color=(0, 0, 0, 255),
                        blur_radius=20,
                    )

                    # 计算在列画布上的位置
                    y_position = row_index * (cell_height + margin)

                    # 粘贴到列画布上
                    column_image.paste(
                        resized_poster_with_shadow,
                        (0, y_position),
                        resized_poster_with_shadow,
                    )

                except Exception as e:
                    logger.error(f"处理图片 {os.path.basename(poster_path)} 时出错: {e}")
                    continue

            # 保存原始列图像
            if save_columns:
                column_orig_path = os.path.join(
                    columns_dir, f"column_{col_index+1}_original.png"
                )
                column_image.save(column_orig_path)
                logger.debug(f"已保存原始列图像到: {column_orig_path}")

            # 旋转整个列
            rotation_canvas_size = int(
                math.sqrt(
                    (cell_width + shadow_extra_width) ** 2
                    + (column_height + shadow_extra_height) ** 2
                )
                * 1.5
            )
            rotation_canvas = Image.new(
                "RGBA", (rotation_canvas_size, rotation_canvas_size), (0, 0, 0, 0)
            )

            # 将列图片放在旋转画布的中央
            paste_x = (rotation_canvas_size - cell_width) // 2
            paste_y = (rotation_canvas_size - column_height) // 2
            rotation_canvas.paste(column_image, (paste_x, paste_y), column_image)

            # 旋转整个列
            rotated_column = rotation_canvas.rotate(
                rotation_angle, Image.BICUBIC, expand=True
            )

            # 保存旋转后的列图像
            if save_columns:
                column_rotated_path = os.path.join(
                    columns_dir, f"column_{col_index+1}_rotated.png"
                )
                rotated_column.save(column_rotated_path)
                logger.debug(f"已保存旋转后的列图像到: {column_rotated_path}")

            # 计算列在模板上的位置
            column_center_y = start_y + column_height // 2
            column_center_x = column_x

            # 根据列索引调整位置
            if col_index == 1:  # 中间列
                column_center_x += cell_width - 50
            elif col_index == 2:  # 右侧列
                column_center_y += -155
                column_center_x += (cell_width) * 2 - 60

            # 计算最终放置位置
            final_x = column_center_x - rotated_column.width // 2 + cell_width // 2
            final_y = column_center_y - rotated_column.height // 2

            # 粘贴旋转后的列到结果图像
            result.paste(rotated_column, (final_x, final_y), rotated_column)

        # 获取第一张图片的随机点颜色，用于色块
        random_color = get_random_color(poster_files[0])

        # 字体路径
        chinese_font_path = os.path.join(os.getcwd(), FONT_FOLDER, CHINESE_FONT)
        english_font_path = os.path.join(os.getcwd(), FONT_FOLDER, ENGLISH_FONT)

        # 添加标题文字
        result = draw_text_on_image(
            result, title, (73.32, 427.34), chinese_font_path, 163,
            shadow_enabled=True, shadow_offset=(2, 2)
        )

        # 添加副标题
        if subtitle:
            # 动态调整字体大小
            base_font_size = 50
            line_spacing = 5

            word_count = len(subtitle.split())
            max_chars_per_line = max([len(word) for word in subtitle.split()])

            if max_chars_per_line > 10 or word_count > 3:
                font_size = base_font_size * (10 / max(max_chars_per_line, word_count * 3)) ** 0.8
                font_size = max(font_size, 30)
            else:
                font_size = base_font_size

            # 使用多行文本绘制
            result, line_count = draw_multiline_text_on_image(
                result, subtitle, (124.68, 624.55), english_font_path,
                int(font_size), line_spacing, shadow_enabled=True, shadow_offset=(2, 2)
            )

            # 根据行数调整色块高度
            color_block_position = (84.38, 620.06)
            color_block_height = 55 + (line_count - 1) * (int(font_size) + line_spacing)
            color_block_size = (21.51, color_block_height)

            result = draw_color_block(
                result, color_block_position, color_block_size, random_color
            )

        # 对整个图像倒圆角
        background_corner_radius = POSTER_GEN_CONFIG["BACKGROUND_CORNER_RADIUS"]
        mask = Image.new("L", (template_width, template_height), 0)
        draw = ImageDraw.Draw(mask)
        draw.rounded_rectangle(
            [(0, 0), (template_width, template_height)],
            radius=background_corner_radius,
            fill=255,
        )
        result = Image.composite(result, Image.new("RGBA", (template_width, template_height), (0, 0, 0, 0)), mask)

        # 如果需要调整分辨率，进行缩放
        if scale_factor != 1.0:
            logger.info(f"正在将海报从1080p缩放到{resolution} ({target_width}x{target_height})")
            result = result.resize(
                (target_width, target_height),
                Image.LANCZOS  # 使用高质量的重采样算法
            )
            
            # 重新应用圆角（因为缩放可能会使圆角边缘变模糊）
            mask = Image.new("L", (target_width, target_height), 0)
            draw = ImageDraw.Draw(mask)
            scaled_corner_radius = int(background_corner_radius * scale_factor)
            draw.rounded_rectangle(
                [(0, 0), (target_width, target_height)],
                radius=scaled_corner_radius,
                fill=255,
            )
            result = Image.composite(result, Image.new("RGBA", (target_width, target_height), (0, 0, 0, 0)), mask)

        # 保存结果
        result.save(output_path, "PNG")
        logger.info(f"成功: 图片已保存到 {output_path}")
        return True

    except Exception as e:
        logger.error(f"创建海报时出错: {e}", exc_info=True)
        return False

if __name__ == "__main__":
    # 检查必要的文件夹是否存在
    if not os.path.exists(TEMP_FOLDER):
        logger.error(f"错误: 找不到temp文件夹，请在当前目录下创建temp文件夹并放入海报图片")
        exit(1)
    
    if not os.path.exists(FONT_FOLDER):
        logger.error(f"错误: 找不到fonts文件夹，请在当前目录下创建fonts文件夹并放入ch.ttf和en.otf字体文件")
        exit(1)
    
    # 检查字体文件是否存在
    chinese_font_path = os.path.join(FONT_FOLDER, CHINESE_FONT)
    english_font_path = os.path.join(FONT_FOLDER, ENGLISH_FONT)
    
    if not os.path.exists(chinese_font_path):
        logger.error(f"错误: 找不到中文字体文件: {chinese_font_path}")
        exit(1)
    
    if not os.path.exists(english_font_path):
        logger.error(f"错误: 找不到英文字体文件: {english_font_path}")
        exit(1)
    
    # 生成海报 - 默认为720p
    generate_poster(
        title="我的电影集", 
        subtitle="My Movie Collection",
        output_filename="poster_output_720p.png",
        resolution="720p"
    )    